# server.py — Manus-like browser-use service
# 兼容方案：优先使用 browser-use 0.6.x 的 cdp-use 客户端；若不可用则回退到 Playwright 截图/输入
import os
import uuid
import base64
import asyncio
from typing import Dict, Any, Optional, List, Callable

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, Header
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from dotenv import load_dotenv

# browser-use & LLM
from browser_use import Agent, BrowserSession
from browser_use.browser import BrowserProfile  # 新版配置入口
from browser_use.llm import ChatDeepSeek

# ---------------- 环境变量 ----------------
load_dotenv()
API_KEY = os.getenv("API_KEY", "dev-key")
OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL", "https://api.deepseek.com/v1")
DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
DEEPSEEK_MODEL = os.getenv("DEEPSEEK_MODEL", "deepseek-chat")

# ---------------- FastAPI & CORS ----------------
app = FastAPI(title="Browser-Use (Manus-like) — CDP优先 + 回退")

# 你的静态资源服务器端口是 5500
ALLOWED_HTTP_ORIGINS = [
    "http://localhost:5500",
    "http://127.0.0.1:5500",
]
app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_HTTP_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

ALLOWED_WS_ORIGINS = set(ALLOWED_HTTP_ORIGINS)


def _check_ws_origin(websocket: WebSocket) -> bool:
    return websocket.headers.get("origin") in ALLOWED_WS_ORIGINS


def check_key(x_api_key: Optional[str]):
    if x_api_key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API key")


# ---------------- Agent 封装 ----------------
class BrowserUseAgent:
    """真实 browser-use Agent，DeepSeek 作为 LLM；支持 pause/resume/cancel。"""

    def __init__(self, browser_session: BrowserSession):
        if not DEEPSEEK_API_KEY:
            raise RuntimeError("Missing DEEPSEEK_API_KEY (set in .env or env vars)")
        self.browser_session = browser_session
        self._paused = False
        self._cancel = False
        print(DEEPSEEK_MODEL)
        print(DEEPSEEK_API_KEY)
        print(OPENAI_BASE_URL)
        self.llm = ChatDeepSeek(
            model=DEEPSEEK_MODEL,
            api_key=DEEPSEEK_API_KEY,
            base_url=OPENAI_BASE_URL,
            temperature=0.8,
        )

    def pause(self):
        self._paused = True

    def resume(self):
        self._paused = False

    def cancel(self):
        self._cancel = True

    async def run(
        self,
        task: str,
        on_step: Optional[Callable[[Dict[str, Any]], None]] = None,
        max_steps: int = 30,
    ) -> Dict[str, Any]:
        agent = Agent(
            task=task,
            llm=self.llm,
            browser_session=self.browser_session,
            max_actions_per_step=6,
        )

        async def _on_step_start(a: Agent):
            try:
                state = await a.browser_session.get_browser_state_summary()
                payload = {
                    "hint": "on_step_start",
                    "url": getattr(state, "url", None),
                    "title": getattr(state, "title", None),
                }
            except Exception:
                payload = {"hint": "on_step_start"}

            if on_step:
                await on_step(payload)

            while self._paused and not self._cancel:
                await asyncio.sleep(0.05)
            if self._cancel:
                raise RuntimeError("Task cancelled")

        async def _on_step_end(a: Agent):
            # 可在此处抓取截图/HTML/数据
            pass

        result = await agent.run(
            on_step_start=_on_step_start,
            on_step_end=_on_step_end,
            max_steps=max_steps,
        )
        return {"ok": True, "summary": "browser-use 任务完成", "raw": str(result)}


# ---------------- 会话对象 ----------------
class Session:
    def __init__(self, sid: str, browser_session: BrowserSession):
        self.id = sid
        self.browser_session = browser_session

        # 兼容两种取法：cdp-use 客户端（优先）、Playwright Page（回退）
        self.cdp = None   # 可能不存在
        self.page = None  # 可能不存在

        self.agent = BrowserUseAgent(browser_session=browser_session)

        self.status = "idle"         # 'idle' | 'running' | 'error'
        self.control_mode = "agent"  # 'agent' | 'human'
        self.operator: Optional[str] = None

        self.event_subs: List[WebSocket] = []
        self.screen_subs: List[WebSocket] = []

        self.lock = asyncio.Lock()
        self.screencast_task: Optional[asyncio.Task] = None


SESSIONS: Dict[str, Session] = {}


# ---------------- 屏幕推流：CDP 优先，Playwright 截图回退 ----------------
async def ensure_page_available(sess: Session):
    """确保 session 有可用的 page 或 cdp 客户端"""
    if sess.cdp is not None or sess.page is not None:
        return True
    
    # 尝试获取 page
    if hasattr(sess.browser_session, "get_page"):
        try:
            page = await sess.browser_session.get_page()
            if page is not None:
                sess.page = page
                return True
        except Exception as e:
            print(f"[ensure_page_available] Failed to get page: {e}")
    
    # 尝试获取 cdp 客户端（优先使用 cdp_client）
    if hasattr(sess.browser_session, "cdp_client"):
        cdp_client = getattr(sess.browser_session, "cdp_client", None)
        if cdp_client is not None:
            sess.cdp = cdp_client
            return True
    
    # 尝试获取 cdp 客户端（旧版本兼容）
    if hasattr(sess.browser_session, "cdp"):
        cdp_client = getattr(sess.browser_session, "cdp", None)
        if cdp_client is not None:
            sess.cdp = cdp_client
            return True
    
    return False

async def ensure_playwright_page(sess: Session) -> bool:
    """仅确保 Playwright page 可用（不因存在 CDP 而返回 True）。"""
    if sess.page is not None:
        return True

    # 尝试通过 browser_session 拉取 page（新版本接口）
    if hasattr(sess.browser_session, "get_page"):
        try:
            page = await sess.browser_session.get_page()
            if page is not None:
                sess.page = page
                return True
        except Exception as e:
            print(f"[ensure_playwright_page] Failed to get page via get_page: {e}")

    # 兼容旧版本直接暴露 page 属性
    page = getattr(sess.browser_session, "page", None)
    if page is not None:
        sess.page = page
        return True

    return False

async def start_screencast(sess: Session, *, quality=70, max_width=1024, max_height=640):
    """
    优先走 cdp-use 的 Page.startScreencast；如果不可用则回退到 page.screenshot() 周期截图。
    """
    # 路径 A：有 cdp-use 客户端
    if sess.cdp is not None and hasattr(sess.cdp, "send") and hasattr(sess.cdp, "register"):
        cdp = sess.cdp
        loop = asyncio.get_running_loop()
        frame_queue: "asyncio.Queue[dict]" = asyncio.Queue()

        def on_frame(evt: dict):
            # evt: {"data": base64, "metadata": {...}, "sessionId": "..."}
            loop.call_soon_threadsafe(frame_queue.put_nowait, evt)

        try:
            cdp.register.Page.screencastFrame(on_frame)
            await cdp.send.Page.enable()
            try:
                await cdp.send.Page.navigate({"url": "about:blank"})
            except Exception:
                pass
            await cdp.send.Page.startScreencast({
                "format": "jpeg",
                "quality": quality,
                "maxWidth": max_width,
                "maxHeight": max_height,
            })
        except Exception as e:
            print("[screencast] CDP start failed, fallback to screenshots:", e)
            try:
                return await _start_screenshot_loop(sess, interval=0.2, quality=quality)
            except Exception as ee:
                print("[screencast] Fallback screenshots cannot start:", ee)
                return None

        async def pump():
            try:
                while True:
                    evt = await frame_queue.get()
                    data_b64 = evt.get("data", "")
                    session_id = evt.get("sessionId")
                    for ws in list(sess.screen_subs):
                        try:
                            await ws.send_text(data_b64)
                        except Exception:
                            try:
                                sess.screen_subs.remove(ws)
                            except Exception:
                                pass
                    if session_id:
                        try:
                            await cdp.send.Page.screencastFrameAck({"sessionId": session_id})
                        except Exception:
                            pass
            finally:
                try:
                    await cdp.send.Page.stopScreencast()
                except Exception:
                    pass

        return asyncio.create_task(pump())

    # 路径 B：Playwright 截图轮询（要求有 page）
    try:
        return await _start_screenshot_loop(sess, interval=0.2, quality=quality)
    except Exception as e:
        print("[screencast] Screenshot loop cannot start:", e)
        return None


async def _start_screenshot_loop(sess: Session, *, interval=0.2, quality=70):
    """
    回退方案：周期性 page.screenshot(type='jpeg', quality=quality)，通过 WS 发 base64。
    """
    page = sess.page
    if page is None:
        # 截图回退必须拿到 Playwright Page，不能因存在 CDP 就放行
        if not await ensure_playwright_page(sess):
            raise RuntimeError("拿不到 Playwright Page，无法进行截图回退推流。")
        page = sess.page

    try:
        await page.goto("about:blank")
    except Exception:
        pass

    async def loop():
        while True:
            try:
                # 若 page 在运行中被置空或失效，尝试重获一次
                current_page = page if page is not None else sess.page
                if current_page is None:
                    if not await ensure_playwright_page(sess):
                        raise RuntimeError("Playwright Page 不可用（重试失败）")
                    current_page = sess.page

                buf: bytes = await current_page.screenshot(type="jpeg", quality=quality)
                data_b64 = base64.b64encode(buf).decode("ascii")
                for ws in list(sess.screen_subs):
                    try:
                        await ws.send_text(data_b64)
                    except Exception:
                        try:
                            sess.screen_subs.remove(ws)
                        except Exception:
                            pass
            except Exception as e:
                print("[screenshot loop] error:", e)
            await asyncio.sleep(interval)

    return asyncio.create_task(loop())


# ---------------- API 模型 ----------------
class CreateSessionResp(BaseModel):
    session_id: str


class TaskReq(BaseModel):
    task: str
    max_steps: int = 30


class TaskStatusResp(BaseModel):
    status: str
    control_mode: str
    operator: Optional[str] = None


# ---------------- 路由：创建 / 删除会话 ----------------
@app.post("/sessions", response_model=CreateSessionResp)
async def create_session(x_api_key: Optional[str] = Header(None)):

    check_key(x_api_key)
    sid = str(uuid.uuid4())

    # keep_alive 正确放在 BrowserProfile
    profile = BrowserProfile(
        keep_alive=True,  # 任务结束不关浏览器，便于复用
        headless=False,   # 开启界面，便于可视化；需要无头可改 True
        # 还可配置：user_data_dir、proxy、viewport 等（按你安装版本的文档）
    )
    browser_session = BrowserSession(browser_profile=profile)
    await browser_session.start()

    sess = Session(sid, browser_session=browser_session)

    # 调试信息（可选）
    # print("browser_use version:", getattr(__import__("browser_use"), "__version__", "unknown"))
    # print("Has cdp_client:", hasattr(browser_session, "cdp_client"))
    # print("Has cdp:", hasattr(browser_session, "cdp"))
    # print("Has page:", hasattr(browser_session, "page"))


    # ① cdp-use 客户端（优先使用 cdp_client）
    cdp_client = getattr(browser_session, "cdp_client", None)
    if cdp_client is None:
        cdp_client = getattr(browser_session, "cdp", None)  # 旧版本兼容
    if cdp_client is not None:
        sess.cdp = cdp_client

    # ② Playwright Page（若版本暴露）
    page = getattr(browser_session, "page", None)
    if page is None and hasattr(browser_session, "get_page"):
        try:
            page = await browser_session.get_page()
        except Exception:
            page = None
    sess.page = page  # 可能仍为 None

    SESSIONS[sid] = sess

    # 不在创建会话时立即启动推流，改为在 /screen WS 连接时按需启动
    sess.screencast_task = None
    
    return CreateSessionResp(session_id=sid)


@app.delete("/sessions/{sid}")
async def delete_session(sid: str, x_api_key: Optional[str] = Header(None)):
    check_key(x_api_key)
    sess = SESSIONS.pop(sid, None)
    if not sess:
        raise HTTPException(404, "Session not found")

    try:
        if sess.screencast_task and not sess.screencast_task.done():
            sess.screencast_task.cancel()
        try:
            await sess.browser_session.stop()
        except Exception:
            pass
    finally:
        return {"ok": True}


# ---------------- 路由：任务提交 / 状态 ----------------
@app.post("/sessions/{sid}/task")
async def run_task(sid: str, req: TaskReq, x_api_key: Optional[str] = Header(None)):
    check_key(x_api_key)
    sess = SESSIONS.get(sid)
    if not sess:
        raise HTTPException(404, "Session not found")
    if sess.status == "running":
        raise HTTPException(409, "A task is already running in this session")

    async with sess.lock:
        sess.status = "running"

        async def on_step(step: Dict[str, Any]):
            for ws in list(sess.event_subs):
                try:
                    await ws.send_json({"type": "step", "payload": step})
                except Exception:
                    try:
                        sess.event_subs.remove(ws)
                    except Exception:
                        pass

        async def runner():
            try:
                # 确保有可用的页面
                if not await ensure_page_available(sess):
                    raise RuntimeError("无法获取浏览器页面，任务无法执行")
                
                # 如果还没有启动推流，现在启动
                if sess.screencast_task is None:
                    sess.screencast_task = await start_screencast(sess)
                
                if sess.control_mode == "human":
                    sess.agent.pause()
                else:
                    sess.agent.resume()
                result = await sess.agent.run(req.task, on_step=on_step, max_steps=req.max_steps)
                for ws in list(sess.event_subs):
                    try:
                        await ws.send_json({"type": "finished", "payload": result})
                    except Exception:
                        pass
                sess.status = "idle"
            except Exception as e:
                sess.status = "error"
                for ws in list(sess.event_subs):
                    try:
                        await ws.send_json({"type": "error", "payload": {"message": str(e)}})
                    except Exception:
                        pass

        asyncio.create_task(runner())
        return {"ok": True, "status": sess.status}


@app.get("/sessions/{sid}/status", response_model=TaskStatusResp)
async def status(sid: str, x_api_key: Optional[str] = Header(None)):
    check_key(x_api_key)
    sess = SESSIONS.get(sid)
    if not sess:
        raise HTTPException(404, "Session not found")
    return TaskStatusResp(
        status=sess.status,
        control_mode=sess.control_mode,
        operator=sess.operator,
    )


# ---------------- WebSocket：事件 / 画面 / 控制 ----------------
@app.websocket("/sessions/{sid}/events")
async def ws_events(websocket: WebSocket, sid: str):
    if not _check_ws_origin(websocket):
        return await websocket.close(code=4403)
    await websocket.accept()
    sess = SESSIONS.get(sid)
    if not sess:
        await websocket.send_json({"type": "error", "payload": {"message": "no session"}})
        return await websocket.close()
    sess.event_subs.append(websocket)
    try:
        while True:
            await websocket.receive_text()  # keepalive
    except WebSocketDisconnect:
        try:
            sess.event_subs.remove(websocket)
        except Exception:
            pass


@app.websocket("/sessions/{sid}/screen")
async def ws_screen(websocket: WebSocket, sid: str):
    if not _check_ws_origin(websocket):
        return await websocket.close(code=4403)
    await websocket.accept()
    sess = SESSIONS.get(sid)
    if not sess:
        return await websocket.close()
    sess.screen_subs.append(websocket)

    # 确保一旦有订阅者接入，就尝试启动推流
    try:
        # 尽力拿到可用的页面或 CDP
        await ensure_page_available(sess)
        # 如果还没有推流任务或任务已结束，则尝试启动
        if sess.screencast_task is None or sess.screencast_task.done():
            task = await start_screencast(sess)
            if task is not None:
                sess.screencast_task = task
            else:
                # 无法启动（例如拿不到 page），先等待客户端后续触发再试
                print("[ws_screen] Screencast not started yet (no page/cdp available)")
    except Exception as e:
        print("[ws_screen] failed to start screencast:", e)
    try:
        while True:
            await websocket.receive_text()  # keepalive
    except WebSocketDisconnect:
        try:
            sess.screen_subs.remove(websocket)
        except Exception:
            pass


@app.websocket("/sessions/{sid}/control")
async def ws_control(websocket: WebSocket, sid: str):
    if not _check_ws_origin(websocket):
        return await websocket.close(code=4403)
    await websocket.accept()
    sess = SESSIONS.get(sid)
    if not sess:
        return await websocket.close()

    user_id = str(uuid.uuid4())  # 生产替换为鉴权后的用户ID
    try:
        while True:
            msg = await websocket.receive_json()
            t = msg.get("type")

            if t == "handover":
                if sess.operator and sess.operator != user_id:
                    await websocket.send_json({"type": "deny", "reason": "occupied"})
                    continue
                
                # 确保有可用的页面
                if not await ensure_page_available(sess):
                    await websocket.send_json({"type": "error", "payload": {"message": "No page available for handover"}})
                    continue
                
                # 如果还没有启动推流，现在启动
                if sess.screencast_task is None:
                    sess.screencast_task = await start_screencast(sess)
                
                sess.operator = user_id
                sess.control_mode = "human"
                sess.agent.pause()
                await websocket.send_json({"type": "granted"})

            elif t == "release":
                if sess.operator == user_id:
                    sess.operator = None
                    sess.control_mode = "agent"
                    sess.agent.resume()
                    await websocket.send_json({"type": "released"})

            elif t == "mouse":
                if sess.operator != user_id:
                    continue
                
                # 确保有可用的页面
                if not await ensure_page_available(sess):
                    await websocket.send_json({"type": "error", "payload": {"message": "No page available"}})
                    continue
                
                name = msg.get("name")  # "mousePressed"/"mouseReleased"/"mouseMoved"/"mouseWheel"
                rel_x = float(msg.get("x", 0))
                rel_y = float(msg.get("y", 0))
                width = float(msg.get("width", 1024))
                height = float(msg.get("height", 640))
                x = rel_x * width
                y = rel_y * height

                # 优先 CDP 注入，否则 Playwright 回退
                if sess.cdp is not None and hasattr(sess.cdp, "send"):
                    payload = {"type": name, "x": x, "y": y}
                    if "button" in msg:
                        payload["button"] = msg["button"]
                    if name == "mouseWheel":
                        payload["deltaY"] = msg.get("deltaY", 0)
                    await sess.cdp.send.Input.dispatchMouseEvent(payload)
                else:
                    if sess.page is None:
                        continue
                    try:
                        if name == "mouseMoved":
                            await sess.page.mouse.move(x, y)
                        elif name == "mousePressed":
                            button = msg.get("button", "left")
                            await sess.page.mouse.move(x, y)
                            await sess.page.mouse.down(button=button)
                        elif name == "mouseReleased":
                            button = msg.get("button", "left")
                            await sess.page.mouse.move(x, y)
                            await sess.page.mouse.up(button=button)
                        elif name == "mouseWheel":
                            delta = int(msg.get("deltaY", 0))
                            await sess.page.mouse.wheel(0, delta)
                    except Exception as e:
                        print("[mouse fallback] error:", e)

            elif t == "key":
                if sess.operator != user_id:
                    continue
                
                # 确保有可用的页面
                if not await ensure_page_available(sess):
                    await websocket.send_json({"type": "error", "payload": {"message": "No page available"}})
                    continue
                
                if sess.cdp is not None and hasattr(sess.cdp, "send"):
                    await sess.cdp.send.Input.dispatchKeyEvent({
                        "type": msg.get("name", "keyDown"),
                        "key": msg.get("key", ""),
                        "code": msg.get("code", "")
                    })
                else:
                    if sess.page is None:
                        continue
                    try:
                        typ = msg.get("name", "keyDown")
                        key = msg.get("key", "")
                        if typ == "keyDown":
                            await sess.page.keyboard.down(key)
                        elif typ == "keyUp":
                            await sess.page.keyboard.up(key)
                        else:
                            await sess.page.keyboard.press(key)
                    except Exception as e:
                        print("[keyboard fallback] error:", e)

    except WebSocketDisconnect:
        if sess.operator == user_id:
            sess.operator = None
            sess.control_mode = "agent"
            sess.agent.resume()


# ---------------- Main ----------------
if __name__ == "__main__":
    import uvicorn
    # 运行示例：
    #   DEEPSEEK_API_KEY=sk-xxx OPENAI_BASE_URL=https://api.deepseek.com/v1 API_KEY=dev-key \
    #   uvicorn server:app --host 0.0.0.0 --port 8000
    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "8000")))
